#include <dialects/Builtin/Builtin.h>
#include <dialects/Func/Func.h>
#include <dialects/PL/PL.h>
#include <gtest/gtest.h>
#include <lcir/ADT/IList.h>
#include <lcir/ADT/ListNode.h>
#include <lcir/ADT/Range.h>
#include <lcir/Dag2Dag/D2DDriver.h>
#include <lcir/Dag2Dag/Pattern.h>
#include <lcir/IR/Argument.h>
#include <lcir/IR/Constant.h>
#include <lcir/IR/Context.h>
#include <lcir/IR/ExprTemps.h>
#include <lcir/IR/Expression.h>
#include <lcir/IR/IRBuilder.h>
#include <lcir/IR/IRRewritter.h>
#include <lcir/IR/Lambda.h>
#include <lcir/IR/LambdaProto.h>
#include <lcir/IR/LambdaType.h>
#include <lcir/IR/Op.h>
#include <lcir/IR/Prototype.h>
#include <lcir/IR/Slot.h>
#include <lcir/IR/SlotCombiner.h>
#include <lcir/IR/Type.h>
#include <lcir/IR/VReg.h>
#include <lcir/IR/ValuePrinter.h>

#include <algorithm>
#include <iterator>
#include <memory>
#include <set>
#include <string>
#include <vector>
#include "lcir/Utils/Result.h"

using namespace lcir;
using namespace lcir::builtin;
using namespace lcir::func;
using namespace lcir::pl;
class TestPattern : public Pattern {
 public:
  // 乘法分配律
  // (a + b) * (a + c) = a * (b + c)
  Result rewrite(Expression *expr, IRRewritter &b) const override {
    std::cout << "开始重写"
              << "\n";
    auto add0 = expr->getOperandValue(0)->getDefExpr();
    auto add1 = expr->getOperandValue(1)->getDefExpr();
    auto add0operands = expr->getOperandValue(0)->getDefExpr()->getOperandValues();
    auto add1operands = expr->getOperandValue(1)->getDefExpr()->getOperandValues();
    auto set0 = std::set(add0operands.begin(), add0operands.end());
    auto set1 = std::set(add1operands.begin(), add1operands.end());
    // 求两个set交集
    std::set<Value *> intersection;
    std::set_intersection(set0.begin(), set0.end(), set1.begin(), set1.end(),
                          std::inserter(intersection, intersection.begin()));
    if (intersection.size() != 1) return failure();
    auto *v = *intersection.begin();
    set0.erase(v);
    set1.erase(v);
    auto *add = b.createAndGetRet<Add>({*set0.begin(), *set1.begin()});
    auto *mul = b.createAndGetRet<Mul>({add, v});
    expr->getResult()->rauw(mul);
    b.replaceValue(expr->getResult(), mul);
    // add0->eraseFromLambda();
    // add1->eraseFromLambda();
    return failure();
  }
};

class D2D : public testing::Test {
 protected:
  void SetUp() override {
    // 初始化lambda
    module = Lambda::get();
    IRBuilder builder(module);
    builder.setInsertPoint(module);
    auto *x = builder.createAndGetRet<DeclVar>(Int32Type::get(), "x");
    auto *y = builder.createAndGetRet<DeclVar>(Int32Type::get(), "y");
    auto *z = builder.createAndGetRet<DeclVar>(Int32Type::get(), "z");
    add1 = builder.createAndGetRet<pl::Add>({x, y});
    add2 = builder.createAndGetRet<pl::Add>({x, z});
    mul = builder.createAndGetRet<Mul>({add1, add2});
    ret = builder.create<Return>({mul});
    // 初始化Pattern
    addNode = std::make_unique<PatternNode>(OpInstance::get<Add>());
    // 初始化一个复杂的Pattern
    complexNode = std::make_unique<PatternNode>(OpInstance::get<Mul>());
    complexNode->setInputsWithInstance({Add::getOpInstance(), Add::getOpInstance()});
    pattern.setRootNode(*complexNode);
  }

  Lambda *module{nullptr};
  Value *add1{nullptr};
  Value *add2{nullptr};
  Value *mul{nullptr};
  Expression *ret{nullptr};
  std::unique_ptr<PatternNode> addNode;
  std::unique_ptr<PatternNode> complexNode;
  TestPattern pattern;
};
// TEST_F(D2D, PatternNode) {
//   std::cout << module->print() << "\n";
//   EXPECT_TRUE(succeeded(addNode->match(add1->getDefExpr())));
//   EXPECT_TRUE(failed(addNode->match(mul->getDefExpr())));
// }

// TEST_F(D2D, Pattern) {
//   std::cout << "加法优化"
//             << "\n";
//   EXPECT_TRUE(failed(pattern.match(add1->getDefExpr())));
//   EXPECT_TRUE(succeeded(pattern.match(mul->getDefExpr())));
//   std::cout << module->print() << "\n";
//   pattern.matchAndRewrite(mul->getDefExpr());
//   std::cout << module->print() << "\n";
// }

TEST_F(D2D, Driver) {
  std::cout << module->print() << "\n";

  for (auto user : add1->getUsers()) {
    user.getExpr()->dump();
  }
  D2DDriver driver(module);
  driver.addPattern(std::make_shared<TestPattern>(pattern));
  driver.run();

  std::cout << module->print() << "\n";
}